import csv
import dataclasses
import datetime
import enum
import matplotlib.axes
import matplotlib.gridspec
import matplotlib.pyplot
import numpy
import pandas
import pandas.core.groupby
import pandas_ta  # 请执行 pandas.DataFrame().ta.indicators() 查看概述,
import pathlib
import scipy.signal
import sqlite3
import talib
from typing import Any, Type, Optional, Union, Callable, Sequence  # , Dict, List, Set, Tuple
import vnpy.trader.database
import vnpy.trader.utility
from vnpy.trader.object import (
    Interval,
    Exchange,
    BarData,
    HistoryRequest,
)
from vnpy.trader.utility import utils, tdx1, tdx2, tdx3, tv, AMAssist


def load_data(
    interval: Interval,
    exchange: Exchange,
    symbol: str,
    start: datetime.datetime,
    end: datetime.datetime,
) -> pandas.DataFrame:
    """"""
    sqlite3_connection: sqlite3.Connection = sqlite3.connect(vnpy.trader.utility.get_file_path("database.db"))

    sql_statement = f"""
    SELECT * FROM dbbardata WHERE interval='{interval.value}' AND exchange='{exchange.value}' AND symbol='{symbol}'
        AND datetime>='{start.strftime("%Y-%m-%d %H:%M:%S")}' AND datetime<='{end.strftime("%Y-%m-%d %H:%M:%S")}'
        ORDER BY datetime ASC;
    """

    df = pandas.read_sql_query(sql=sql_statement, con=sqlite3_connection)

    df.rename(columns={'open_price': 'open', 'high_price': 'high', 'low_price': 'low', 'close_price': 'close'}, inplace=True)

    df["datetime"] = pandas.to_datetime(df["datetime"])  # 字符串转时间戳

    # inplace: 是否在原 DataFrame 上修改, 默认为 False
    # verify_integrity: 是否检查索引有无重复, 默认为 False, 若设置为 True 会影响程序性能,
    df.set_index("datetime", inplace=True, verify_integrity=False)

    return df


def load_data_deprecated(
    interval: Interval,
    exchange: Exchange,
    symbol: str,
    start: datetime.datetime,
    end: datetime.datetime,
) -> pandas.DataFrame:
    """"""
    req: HistoryRequest = HistoryRequest(
        symbol=symbol,
        exchange=exchange,
        start=start,
        end=end,
        interval=interval,
    )

    database: vnpy.trader.database.BaseDatabase = vnpy.trader.database.get_database()

    bars: list[BarData] = database.load_bar_data(
        symbol=req.symbol,
        exchange=req.exchange,
        interval=req.interval,
        start=req.start,
        end=req.end,
    )

    df: pandas.DataFrame = pandas.DataFrame()

    if bars:
        for bar in bars:
            bar.datetime = bar.datetime.replace(tzinfo=None)
            bar.exchange = bar.exchange.value
            bar.interval = bar.interval.value

        df: pandas.DataFrame = pandas.DataFrame(data=[bar.__dict__ for bar in bars])

        df.drop(columns="gateway_name", inplace=True)

        df.rename(columns={'open_price': 'open', 'high_price': 'high', 'low_price': 'low', 'close_price': 'close'}, inplace=True)

        # inplace: 是否在原 DataFrame 上修改, 默认为 False
        # verify_integrity: 是否检查索引有无重复, 默认为 False, 若设置为 True 会影响程序性能,
        df.set_index("datetime", inplace=True, verify_integrity=False)

    return df


def view_group(pandas_group: Union[pandas.core.groupby.SeriesGroupBy, pandas.core.groupby.DataFrameGroupBy]):
    """
    https://blog.csdn.net/m0_46478042/article/details/137439509
    """
    assert isinstance(pandas_group, (pandas.core.groupby.SeriesGroupBy, pandas.core.groupby.DataFrameGroupBy))

    for name, group in pandas_group:
        assert isinstance(group, (pandas.DataFrame, pandas.Series))

        cols = group.name if isinstance(group, pandas.Series) else group.columns.to_list()
        print(f'name={name}, cols={cols}')

        print('-' * 30)
        print(group)
        print('=' * 30, '\n')


class PltHelper(object):
    """"""

    class Color(enum.Enum):
        """"""
        # 在 matplotlib.axes.Axes.plot 的帮助文档搜索『**Colors**』
        blue = 'b'  # 也可以这样写"#0000FF"
        green = 'g'  # 也可以这样写"#00FF00"
        red = 'r'  # 也可以这样写"#FF0000"
        cyan = 'c'
        magenta = 'm'
        yellow = 'y'
        black = 'k'
        white = 'w'
        Orange = "#FFA500"  # 橙色

    class Cate(enum.Enum):
        PLOT = "plot"
        SCATTER = "scatter"

    @dataclasses.dataclass
    class DataItem:
        """"""
        series: pandas.Series
        cate: "PltHelper.Cate"
        size: float | None
        color: str | None

    @dataclasses.dataclass
    class PlotItem:
        """"""
        title: str
        height: int
        data_list: list["PltHelper.DataItem"] = dataclasses.field(default_factory=list)

    def __init__(self, sharex: bool = False) -> None:
        self._sharex: bool = sharex
        self._cache: dict[str, PltHelper.PlotItem] = {}

    def add_plot(self, title: str, height: int = 1) -> None:
        """"""
        item = PltHelper.PlotItem(title=title, height=height)
        self._cache[item.title] = item

    def add_data(
        self,
        title: str,
        series: pandas.Series,
        cate: Cate = Cate.PLOT,
        size: float | None = None,
        color: str | Color | None = None,
    ) -> None:
        """"""
        item = PltHelper.DataItem(
            series=series,
            cate=cate,
            size=size,
            color=color.value if isinstance(color, PltHelper.Color) else color,
        )
        self._cache[title].data_list.append(item)

    def show(self):
        """"""
        # python如何调整子图的大小？
        # https://www.zhihu.com/question/21953954

        nrows: int = sum([_.height for _ in self._cache.values()])

        fig = matplotlib.pyplot.figure()
        gs = matplotlib.gridspec.GridSpec(nrows=nrows, ncols=1, figure=fig)

        if self._sharex:
            fig.subplots_adjust(hspace=0)  # 为了更好的可视化需求, 可以把子区间的空隙去掉.

        sharex: Optional[matplotlib.axes.Axes] = None

        idx: int = 0
        for row, item in enumerate(self._cache.values()):
            beg, end, idx = idx, (idx + item.height), (idx + item.height)
            ax: matplotlib.axes.Axes = fig.add_subplot(gs[beg:end, 0], sharex=sharex)

            sharex = ax if self._sharex and (sharex is None) else sharex

            if self._sharex and (len(self._cache) != row + 1):
                matplotlib.pyplot.setp(ax.get_xticklabels(), visible=False)

            for data in item.data_list:
                data: PltHelper.DataItem = data

                match data.cate:
                    case PltHelper.Cate.PLOT:
                        ax.plot(data.series.index, data.series, label=data.series.name, markersize=data.size, color=data.color)
                    case PltHelper.Cate.SCATTER:
                        # https://blog.csdn.net/Poul_henry/article/details/88602806
                        # scatter()中参数s表示的大小，和plot()的markersize大小的平方相同，即[s]=markersize**2
                        ax.scatter(x=data.series.index, y=data.series, label=data.series.name, s=data.size, c=data.color)
                    case _:
                        raise NotImplementedError

            ax.set_title(item.title)
            ax.legend()  # 添加图例/图注
            ax.grid()  # 添加网格

        matplotlib.pyplot.show()  # 展示图像

    def show_deprecated(self):
        """"""
        # 在创建子图时，可以通过设置 constrained_layout=True 参数来激活约束布局。这个参数会自动调整子图和装饰，使它们尽可能地适合画布大小。
        # sharex, sharey：是否共享x/y轴
        # 共享x轴之后，x轴坐标只显示在最下面的图下
        fig, axes = matplotlib.pyplot.subplots(nrows=len(self._cache), ncols=1, sharex=self._sharex)

        # 在 Python 中使用共享轴的 GridSpec
        # https://segmentfault.com/q/1010000043245355
        if self._sharex:
            fig.subplots_adjust(hspace=0)  # 为了更好的可视化需求, 可以把子区间的空隙去掉.

        for row, item in enumerate(self._cache.values()):
            ax: matplotlib.axes.Axes = axes[row]

            for data in item.data_list:
                ax.plot(data.series.index, data.series, label=data.series.name)

            ax.set_title(item.title)
            ax.legend()  # 添加图例/图注
            ax.grid()  # 添加网格

        matplotlib.pyplot.show()  # 展示图像


def calculate_statistics(
    df: pandas.DataFrame,
    output: bool = True,
    show_trade: bool = False,
    show_graph: bool = False,
    rate: float = 0,  # 手续费率,
    slippage: float = 0,  # 交易滑点(这里的滑点的单位:pricetick的单位, 传统的滑点的单位:pricetick, 这里的滑点=传统的滑点*pricetick)
    size: float = 1,  # 合约乘数,
    capital: float = 0,  # 初始资金,
) -> None:
    """"""
    data: pandas.DataFrame = df.copy()

    if show_trade:
        # 如果我们粗暴地设定"成交价"是K线的收盘价, 那么, 成交数据如下所示:
        print(data[data["position"] != data["position"].shift(1)])

    if 0:
        # 计算收益
        data['returns'] = data['close'].pct_change()  # 当日收益率
        data['strategy_returns'] = data['returns'] * data['position'].shift(1)  # 策略当日收益(使用前一日仓位)
        # 计算累积收益
        data['cumulative_strategy'] = (1 + data['strategy_returns']).cumprod()
        data['cumulative_buy_hold'] = (1 + data['returns']).cumprod()

    # 计算手续费
    data["pos_difabs"] = data["position"].diff().abs().fillna(0)
    data["slippage"] = data["pos_difabs"] * slippage * size  # 滑点(这里的 slippage = 多少个"价格跳动"的滑点 * 价格跳动)
    data["turnover"] = data["pos_difabs"] * data["close"] * size  # 成交额
    data["commission"] = data["turnover"] * rate  # 手续费
    # 计算当日收益
    data["px_diff"] = data["close"].diff()
    data["net_pnl"] = data["px_diff"] * data["position"].shift(1) * size - data["commission"] - data["slippage"]  # 策略当日收益(绝对金额)
    # 计算累积收益
    data["strategy_cum"] = capital + data["net_pnl"].cumsum()  # 策略累积收益(绝对金额)
    data["buy_hold_cum"] = capital + data["px_diff"].cumsum() * 1 * size  # Buy&Hold累积收益(绝对金额, 设定初始买入1手)
    # 计算回撤
    data["drawdown"] = data["strategy_cum"] - data["strategy_cum"].cummax()  # 回撤(绝对金额, 暂不计算相对回撤)

    if output:
        print("起始资金:", "{:,.6f}".format(capital))
        print("结束资金:", "{:,.6f}".format(data.iloc[-1]["strategy_cum"]))
        print("总盈亏：:", "{:,.6f}".format(data["net_pnl"].sum()))
        print("总手续费:", "{:,.6f}".format(data["commission"].sum()))
        print("总滑点：:", "{:,.6f}".format(data["slippage"].sum()))
        print("成交金额:", "{:,.6f}".format(data["turnover"].sum()))
        print("成交笔数:", "{:,.6f}".format(data["pos_difabs"].sum()))

    if show_graph:
        df_p = data.reset_index(inplace=False)
        helper = PltHelper(sharex=False)
        helper.add_plot("candle", 2)
        helper.add_plot("position", 1)
        helper.add_plot("net_pnl", 1)
        helper.add_plot("strategy_cum", 1)
        helper.add_plot("drawdown", 1)
        helper.add_data("candle", df_p["close"], cate=PltHelper.Cate.PLOT)
        helper.add_data("position", df_p["position"], cate=PltHelper.Cate.PLOT)
        helper.add_data("net_pnl", df_p["net_pnl"], cate=PltHelper.Cate.PLOT)
        helper.add_data("strategy_cum", df_p["strategy_cum"], cate=PltHelper.Cate.PLOT)
        helper.add_data("drawdown", df_p["drawdown"], cate=PltHelper.Cate.PLOT)
        helper.show()

    return


def example():
    """"""
    dd_beg = datetime.date(year=2024, month=1, day=1)
    dd_end = datetime.date(year=2025, month=2, day=21)

    dt_beg = datetime.datetime.combine(date=dd_beg, time=datetime.time(hour=19)) - datetime.timedelta(days={0: 3, 5: 1, 6: 2}.get(dd_beg.weekday(), 1))
    dt_end = datetime.datetime.combine(date=dd_end, time=datetime.time(hour=17))

    df = load_data(
        exchange=Exchange.SHFE,
        symbol="agL8",
        interval=Interval.MINUTE,
        start=dt_beg - datetime.timedelta(days=30),
        end=dt_end,
    )

    dt_1 = datetime.datetime.now()

    df["ma_fast"] = df['close'].rolling(38).mean()
    df["ma_slow"] = df['close'].rolling(62).mean()

    df["signal"] = numpy.nan
    df.loc[(df["ma_fast"] > df["ma_slow"]) & (df["ma_fast"] <= df["ma_slow"]).shift(1), "signal"] = +1
    df.loc[(df["ma_fast"] < df["ma_slow"]) & (df["ma_fast"] >= df["ma_slow"]).shift(1), "signal"] = -1

    df = df.loc[dt_beg <= df.index]  # 丢弃前置数据,

    # 计算仓位
    df["position"] = numpy.nan
    df.loc[df['signal'] == +1, "position"] = +1
    df.loc[df['signal'] == +0, "position"] = +0
    df.loc[df['signal'] == -1, "position"] = -1
    # 为什么要『shift(1)』只有这一根K线完成了, 才能计算这一根K线的"signal",
    # 如果要下单, 最早也得在下一根K线期间下单, 进而改变"position"的值,
    df["position"] = df["position"].shift(1).ffill().fillna(0)  # 前向填充仓位

    calculate_statistics(
        df=df,
        output=True,
        show_trade=False,
        show_graph=False,
        rate=5.007e-05,
        slippage=0.000667,
        size=15.0,
        capital=100000.0,
    )

    dt_2 = datetime.datetime.now()

    print("向量化计算耗时:", (dt_2 - dt_1).total_seconds())


def save_trade_record(data: pandas.DataFrame, sign: int) -> None:
    """"""
    # awk如果指定多个分隔符, 使用中括号[]
    # awk -F '[,+]' -v OFS="," '{print $9,$3,$4,$5,$6,$7}'  "trade_record.csv" > tp.csv
    lines: list[dict] = []

    match sign:
        case 1:
            bs: str = "多空"
        case -1:
            bs: str = "空多"
        case _:
            raise NotImplementedError

    for idx, row in data.iterrows():
        line = {"成交号": "_", "委托号": "_", "代码": row["symbol"], "交易所": row["exchange"], "方向": bs[0], "开平": "开", "价格": row["px_i"], "数量": "_", "时间": row["dt_i"], "接口": "_"}
        lines.append(line)
        line = {"成交号": "_", "委托号": "_", "代码": row["symbol"], "交易所": row["exchange"], "方向": bs[1], "开平": "平", "价格": row["px_o"], "数量": "_", "时间": row["dt_o"], "接口": "_"}
        lines.append(line)

    with open(file=pathlib.Path(__file__).parent / "trade_record.csv", mode="w", encoding="utf8", newline='') as fp:
        if lines:
            csvWriter = csv.DictWriter(f=fp, fieldnames=lines[0].keys())
            csvWriter.writeheader()
            csvWriter.writerows(lines)


def do_describe(df: pandas.DataFrame, col_name: str, bar_num: int, sign: int, vnpy: bool | int = 0):
    """立即下单, 并持仓 bar_num 根K线, 的价格的盈亏情况"""
    data: pandas.DataFrame = df.copy()
    assert sign in (-1, +1)
    assert 0 < bar_num  # n根K线之后, 价格的差异,

    if vnpy:
        data["dttm"] = data.index
        # 假设总是以开盘价成交, 手续费率=0, 交易滑点=0, 合约乘数=1, 下面这样计算, 数据和 vnpy 的数据(基本)一致,
        # 另: 需要修改 vnpy_ctastrategy/backtesting.py 的 cross_limit_order 函数, 使其按开盘价成交,
        data["px_i"] = data["open"].shift(-1)
        data["px_o"] = data["open"].shift(+0 - bar_num).shift(-1)
        data["dt_i"] = data["dttm"].shift(-1)
        data["dt_o"] = data["dttm"].shift(+0 - bar_num).shift(-1)
        data["diff_tmp"] = (data["px_o"] - data["px_i"]) * sign
        data["diff_val"] = (data["open"].shift(+0 - bar_num) - data["open"]).shift(-1) * sign
        assert data["diff_val"].equals(data["diff_tmp"])
        # 如果这里多了几条数据, 这几条数据, 按字面上的多空开平计算出来了一个(含正负号的)数字,
        # 其实这个数字已经叠加到 sum 里了,所以要减掉这个数字, 才能和 vnpy 的数据做对比,
        save_trade_record(data=data[data[col_name] == 1], sign=sign)
    else:
        # 只有这一根K线完成了, 才能计算这一根K线的 col_name 的值,
        # 如果要下单, 最早也得在下一根K线期间下单, 假设下单非常及时, 在"open"的地方成交了,
        data["diff_val"] = (data["close"].shift(1 - bar_num) - data["open"]).shift(-1) * sign

    print(f"列名: {col_name}, 持仓 {bar_num} 根K线, 符号 {sign:+d}:")
    print(data[data[col_name] == 1]["diff_val"].describe())
    print("sum:", data[data[col_name] == 1]["diff_val"].sum())
    print()


def merge_macd(df: pandas.DataFrame, df_macd: pandas.DataFrame):
    """"""
    # 例如 df_macd = pandas_ta.macd(close=df["close"], fast=12, slow=26, signal=9)
    mapping = {"MACD_": "dif", "MACDs_": "dea", "MACDh_": "macd"}
    for column in df_macd.columns:
        key: str = column.split(sep='_', maxsplit=1)[0] + '_'
        df[mapping[key]] = df_macd[column]


def example_2():
    """"""
    dd_beg = datetime.date(year=2024, month=1, day=2)
    dd_end = datetime.date(year=2025, month=3, day=4)

    dt_beg = datetime.datetime.combine(date=dd_beg, time=datetime.time(hour=19)) - datetime.timedelta(days={0: 3, 5: 1, 6: 2}.get(dd_beg.weekday(), 1))
    dt_end = datetime.datetime.combine(date=dd_end, time=datetime.time(hour=17))

    df = load_data(
        exchange=Exchange.SHFE,
        symbol="agL8",
        interval=Interval.MINUTE05,
        start=dt_beg - datetime.timedelta(days=30),
        end=dt_end,
    )

    df["ma_fast"] = tdx1.MA(df['close'], 38)
    df["ma_slow"] = tdx1.MA(df['close'], 62)

    df["cross_up"] = tv.CrossOver(df["ma_fast"], df["ma_slow"], array=True)
    df["cross_dn"] = tv.CrossUnder(df["ma_fast"], df["ma_slow"], array=True)

    df = df.loc[dt_beg <= df.index]  # 丢弃前置数据,

    do_describe(df, "cross_up", bar_num=5, sign=+1, vnpy=0)
    do_describe(df, "cross_dn", bar_num=5, sign=-1, vnpy=0)


if __name__ == "__main__":
    if 0:
        pandas.set_option('display.max_rows', None)
    example_2()
